import pandas as pd
import numpy as np
import requests
from bs4 import BeautifulSoup
import time
import re
import urllib.parse
import os
import random
import traceback
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# Import functionality from the PubChem API script
from API_PubChem_New import search_pubchem_by_cas as search_pubchem

def create_session_with_retry():
    """
    Create a session with retry capabilities
    """
    session = requests.Session()
    
    # Configure retry strategy
    retries = Retry(
        total=5,
        backoff_factor=1,
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["GET", "POST"]
    )
    
    # Mount the adapter to the session
    adapter = HTTPAdapter(max_retries=retries)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    
    return session

def get_browser_headers():
    """
    Return realistic browser headers
    """
    return {
        'User-Agent': 'Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 Mobile/15E148 Safari/604.1',
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
        'Accept-Language': 'zh-CN,zh;q=0.9',
        'Connection': 'keep-alive'
    }

def extract_chemical_info(html_content):
    """
    Extract chemical information from the boxList div
    """
    # Initialize properties dictionary
    properties = {
        "中文名称": "",
        "英文名称": "",
        "分子式": "",
        "CAS号": "",
        "活性": "",
        "药理作用": "",
        "CB号": ""
    }
    
    # Parse HTML
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # Find the boxList div
    box_list = soup.find('div', id='boxList')
    if not box_list:
        print("boxList div not found")
        return properties
    
    # Extract list items
    list_items = box_list.find_all('li')
    
    for item in list_items:
        # Extract raw text and remove extra whitespace
        item_text = item.get_text(strip=True)
        
        # Extract CBNumber
        cb_link = item.select_one('a[href*="ProductChemicalPropertiesCB"]')
        if cb_link:
            cb_match = re.search(r'ProductChemicalPropertiesCB(\d+)', cb_link['href'])
            if cb_match:
                properties["CB号"] = f"CB{cb_match.group(1)}"
        
        # Extract English name
        if '英文名称' in item_text:
            # The English name is the text after '英文名称'
            english_name = item_text.replace('英文名称', '', 1).strip()
            properties["英文名称"] = english_name
        
        # Extract Chinese name
        if '中文名称' in item_text:
            # The Chinese name is the text after '中文名称'
            chinese_name = item_text.replace('中文名称', '', 1).strip()
            properties["中文名称"] = chinese_name
        
        # Extract molecular formula (MF)
        if 'MF' in item_text:
            # The formula is the text after 'MF'
            formula = item_text.replace('MF', '', 1).strip()
            properties["分子式"] = formula
            
        # Extract CAS number
        if item_text.startswith('CAS'):
            # The CAS number is the text after 'CAS'
            cas = item_text.replace('CAS', '', 1).strip()
            properties["CAS号"] = cas
    
    return properties

def get_additional_chemical_details(session, cb_number):
    """
    Get additional chemical details from the CB number URL
    """
    # Default empty additional properties
    additional_properties = {
        "活性": "",
        "药理作用": ""
    }
    
    if not cb_number:
        return additional_properties
    
    # Construct the full URL
    details_url = f"https://m.chemicalbook.com/ProductChemicalPropertiesCB{cb_number}.htm"
    
    try:
        # Add random delay 
        time.sleep(random.uniform(2, 5))
        
        # Send request
        response = session.get(details_url, headers=get_browser_headers(), timeout=30)
        
        # Check response
        if response.status_code != 200:
            print(f"Failed to fetch details for {cb_number}: HTTP {response.status_code}")
            return additional_properties
        
        # Parse HTML
        soup = BeautifulSoup(response.text, 'html.parser')
        
        # Extract biological activity
        biological_activity_div = soup.find('div', class_='gsbt', string=re.compile('生物活性'))
        if biological_activity_div and biological_activity_div.next_sibling:
            biological_activity = biological_activity_div.next_sibling.strip()
            additional_properties["活性"] = biological_activity
        
        # Extract pharmacological effects
        pharmacological_effects_div = soup.find('div', class_='gsbt', string=re.compile('用途'))
        if pharmacological_effects_div and pharmacological_effects_div.next_sibling:
            pharmacological_effects = pharmacological_effects_div.next_sibling.strip()
            additional_properties["药理作用"] = pharmacological_effects
        
        return additional_properties
    
    except Exception as e:
        print(f"Error fetching details for {cb_number}: {e}")
        return additional_properties

def search_chemicalbook_by_cas(session, cas_number):
    """
    Search ChemicalBook using CAS number
    """
    # Initialize an empty properties dictionary
    empty_properties = {
        "中文名称": "",
        "英文名称": "",
        "分子式": "",
        "CAS号": "",
        "活性": "",
        "药理作用": "",
        "CB号": ""
    }
    
    if not cas_number or pd.isna(cas_number) or cas_number == "":
        return empty_properties
    
    # Clean the CAS number
    cas_number = str(cas_number).strip()
    
    # Encode the CAS number for the URL
    encoded_cas = urllib.parse.quote(cas_number)
    search_url = f"https://m.chemicalbook.com/Search.aspx?keyword={encoded_cas}"
    
    try:
        # Add random delay to mimic human behavior
        time.sleep(random.uniform(2, 5))
        
        # Send request
        response = session.get(search_url, headers=get_browser_headers(), timeout=30)
        
        # Check response
        if response.status_code != 200:
            print(f"Failed to search for CAS '{cas_number}': HTTP {response.status_code}")
            return empty_properties
        
        # Extract chemical information
        properties = extract_chemical_info(response.text)
        
        # Verify the CAS number matches
        if properties["CAS号"] and properties["CAS号"].strip() != cas_number:
            print(f"CAS number mismatch: searched for '{cas_number}', but found '{properties['CAS号']}'")
            return empty_properties
            
        return properties
    
    except Exception as e:
        print(f"Error searching for CAS '{cas_number}': {e}")
        return empty_properties

def search_chemicalbook_by_chinese_name(session, chinese_name):
    """
    Search ChemicalBook using Chinese name
    """
    # Initialize an empty properties dictionary
    empty_properties = {
        "中文名称": "",
        "英文名称": "",
        "分子式": "",
        "CAS号": "",
        "活性": "",
        "药理作用": "",
        "CB号": ""
    }
    
    if not chinese_name or pd.isna(chinese_name) or chinese_name == "" or chinese_name == "随便":
        return empty_properties
    
    # Clean the Chinese name
    chinese_name = str(chinese_name).strip()
    
    # Encode the Chinese name for the URL
    encoded_name = urllib.parse.quote(chinese_name)
    search_url = f"https://m.chemicalbook.com/Search.aspx?keyword={encoded_name}"
    
    try:
        # Add random delay to mimic human behavior
        time.sleep(random.uniform(2, 5))
        
        # Send request
        response = session.get(search_url, headers=get_browser_headers(), timeout=30)
        
        # Check response
        if response.status_code != 200:
            print(f"Failed to search for Chinese name '{chinese_name}': HTTP {response.status_code}")
            return empty_properties
        
        # Extract chemical information
        properties = extract_chemical_info(response.text)
        
        # Report mismatch but continue processing (CHANGED: don't return empty properties)
        if properties["中文名称"] and not (
            chinese_name.lower() in properties["中文名称"].lower() or 
            properties["中文名称"].lower() in chinese_name.lower()
        ):
            print(f"Chinese name mismatch: searched for '{chinese_name}', but found '{properties['中文名称']}'")
            # Continue instead of returning empty properties
            
        return properties
    
    except Exception as e:
        print(f"Error searching for Chinese name '{chinese_name}': {e}")
        return empty_properties

def get_pubchem_data_by_cas(cas_number):
    """
    Get additional data from PubChem specifically using CAS number
    """
    # Skip if no valid CAS number
    if not cas_number or pd.isna(cas_number) or cas_number == "" or cas_number == "9999-9999-9999":
        return {}
    
    try:
        print(f"Searching PubChem for CAS: {cas_number}")
        pubchem_data = search_pubchem(cas_number)
        return pubchem_data
    except Exception as e:
        print(f"Error searching PubChem for CAS '{cas_number}': {e}")
        return {}

def process_excel_file(input_file, output_file):
    """
    Process the input Excel file:
    1. First try searching by CAS number
    2. If no match or CAS is missing, try searching by Chinese name
    3. Fill in missing data from search results - only filling blank cells
    4. If a new CAS number is inserted, immediately search PubChem for CID and SMILES
    5. Add a "随便" row at the beginning and remove it before saving
    6. Add "Yes" to "CB使用中文名" column when Chinese name search is successful
    """
    # Read the input Excel file
    df = pd.read_excel(input_file)
    
    # Check if the required columns exist
    required_columns = ["CAS号", "中文名"]
    missing_columns = [col for col in required_columns if col not in df.columns]
    if missing_columns:
        print(f"Error: Required columns {', '.join(missing_columns)} not found in the Excel file.")
        print(f"Available columns: {', '.join(df.columns)}")
        return
    
    # Add the "CB使用中文名" column if it doesn't exist
    if "CB使用中文名" not in df.columns:
        df["CB使用中文名"] = ""
    
    # Add the "随便" row at the beginning
    print("Adding '随便' row at the beginning for API compatibility")
    new_row = pd.DataFrame({
        col: ['随便' if col == '中文名' else 
              '9999-9999-9999' if col == 'CAS号' else 
              None] 
        for col in df.columns
    }, index=[0])
    df = pd.concat([new_row, df]).reset_index(drop=True)
    
    # Create a session for all requests
    session = create_session_with_retry()
    
    # Iterate through each row in the input DataFrame
    for index, row in df.iterrows():
        cas_number = row.get("CAS号", "")
        chinese_name = row.get("中文名", "")
        english_name = row.get("英文名", "")
        
        # Don't skip the first row with "随便" (process it like any other row)
        
        properties = None
        search_method = ""
        cas_was_inserted = False
        chinese_name_search_successful = False
        
        # First try searching by CAS number if available
        if not pd.isna(cas_number) and cas_number != "":
            print(f"Processing {index+1}/{len(df)}: CAS={cas_number}")
            properties = search_chemicalbook_by_cas(session, cas_number)
            if properties["CAS号"]:
                search_method = "CAS"
        
        # If CAS search failed or CAS is missing, try searching by Chinese name
        if (properties is None or not properties["CAS号"]) and not pd.isna(chinese_name) and chinese_name != "":
            print(f"Processing {index+1}/{len(df)}: Chinese name={chinese_name}")
            properties = search_chemicalbook_by_chinese_name(session, chinese_name)
            if properties["中文名称"] or properties["CAS号"]:  # CHANGED: Accept result even with name mismatch
                search_method = "Chinese name"
                chinese_name_search_successful = True
                # Mark that Chinese name was used successfully
                df.at[index, "CB使用中文名"] = "Yes"
        
        # Skip if no match found (empty properties)
        if properties is None or (not properties["CAS号"] and not properties["中文名称"]):
            print(f"Skipping row {index+1}: No match found")
            continue
            
        print(f"Found match using {search_method}")
        
        # If CB号 exists, fetch additional details
        if properties["CB号"]:
            cb_number = properties["CB号"].replace("CB", "")
            additional_details = get_additional_chemical_details(session, cb_number)
            
            # Merge the additional details with existing properties
            properties.update(additional_details)
        
        # Update Chinese name ONLY if it's blank in the original data
        if properties["中文名称"] and "中文名" in df.columns:
            current_cn_name = row.get("中文名", "")
            if pd.isna(current_cn_name) or current_cn_name == "":
                df.at[index, "中文名"] = properties["中文名称"]
                print(f"  Added Chinese name: {properties['中文名称']}")
        
        # Update English name ONLY if it's blank
        if properties["英文名称"] and "英文名" in df.columns:
            current_en_name = row.get("英文名", "")
            if pd.isna(current_en_name) or current_en_name == "":
                df.at[index, "英文名"] = properties["英文名称"]
                print(f"  Added English name: {properties['英文名称']}")
            
        # Update CAS number if missing or empty
        if properties["CAS号"] and "CAS号" in df.columns:
            current_cas = row.get("CAS号", "")
            if pd.isna(current_cas) or current_cas == "":
                df.at[index, "CAS号"] = properties["CAS号"]
                print(f"  Added CAS number: {properties['CAS号']}")
                cas_was_inserted = True
        
        # Update other fields
        if "活性" in df.columns and properties["活性"]:
            df.at[index, "活性"] = properties["活性"]
            
        if "药理作用" in df.columns and properties["药理作用"]:
            df.at[index, "药理作用"] = properties["药理作用"]
            
        if "CB号" in df.columns and properties["CB号"]:
            df.at[index, "CB号"] = properties["CB号"]
        
        # Update molecular formula if available
        if "分子式" in df.columns and properties["分子式"]:
            df.at[index, "分子式"] = properties["分子式"]
        
        # If a new CAS number was inserted, immediately search PubChem for additional data
        if cas_was_inserted:
            # Get the newly inserted CAS number
            new_cas = properties["CAS号"]
            
            # Check if CID or SMILES is missing
            cid_missing = "CID号" in df.columns and (pd.isna(row.get("CID号", "")) or row.get("CID号", "") == "")
            smiles_missing = "SMILES" in df.columns and (pd.isna(row.get("SMILES", "")) or row.get("SMILES", "") == "")
            
            if cid_missing or smiles_missing:
                # Search PubChem using the CAS number
                pubchem_data = get_pubchem_data_by_cas(new_cas)
                
                # Update CID if missing
                if cid_missing and pubchem_data.get("CID", ""):
                    df.at[index, "CID号"] = pubchem_data["CID"]
                    print(f"  Added CID from PubChem: {pubchem_data['CID']}")
                
                # Update SMILES if missing
                if smiles_missing and pubchem_data.get("SMILES", ""):
                    df.at[index, "SMILES"] = pubchem_data["SMILES"]
                    print(f"  Added SMILES from PubChem")
                
                # Update molecular formula if missing
                if "分子式" in df.columns and pubchem_data.get("MF", "") and (pd.isna(df.at[index, "分子式"]) or df.at[index, "分子式"] == ""):
                    df.at[index, "分子式"] = pubchem_data["MF"]
                    print(f"  Added molecular formula from PubChem: {pubchem_data['MF']}")
    
    # Remove the "随便" row before saving
    if len(df) > 0 and df.iloc[0]['中文名'] == '随便':
        print("Removing the '随便' row before saving")
        df = df.iloc[1:].reset_index(drop=True)
    
    # Save the results to the output Excel file
    df.to_excel(output_file, index=False)
    print(f"Results saved to {output_file}")

def main():
    input_file = "./API/mjz0320_matched_PC.xlsx"
    output_file = "./API/mjz0320_ChemicalBook_结果.xlsx"
    
    # Check if the input file exists
    if not os.path.exists(input_file):
        print(f"Input file '{input_file}' not found.")
        return
    else:
        print(f"Found input file: {input_file}")
    
    # Process the file
    process_excel_file(input_file, output_file)

if __name__ == "__main__":
    main()